package com.su02.multi.examples.mlp;

import com.su02.multi.chainrule.Constant;
import com.su02.multi.chainrule.Function;
import com.su02.multi.chainrule.Variable;

import org.apache.commons.collections4.CollectionUtils;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;

public class Linear4MLP extends Function {
    private List<Variable> theta;
    private Function f;

    private Linear4MLP() {}

    @Override
    public Function doBackward(int id) {
        return f.backward(id);
    }

    @Override
    public double forward(Map<Integer, Double> vector) {
        return f.forward(vector);
    }

    public static Linear4MLP fromVector(List<Double> x, Iterator<Integer> idGenerator) {
        return fromFunctions(x.stream().map(Constant::new).toList(), idGenerator);
    }

    public static Linear4MLP fromFunctions(List<? extends Function> fs, Iterator<Integer> idGenerator) {
        Linear4MLP ans = new Linear4MLP();
        int nx = CollectionUtils.size(fs);
        int nTheta = nx + 1;
        ans.theta = new ArrayList<>(nTheta);
        for (int i = 0; i < nTheta; i++) {
            ans.theta.add(new Variable(tryGetNextId(idGenerator)));
        }
        // entry of θ_i * x_i
        Function f = new Constant(0.);
        for (int i = 0; i < nx; i++) {
            f = f.add(ans.theta.get(i).mul(fs.get(i)));
        }
        f = f.add(ans.theta.get(nTheta - 1));
        ans.f = f;
        ans.setBs(f.getBs());
        return ans;
    }

    private static int tryGetNextId(Iterator<Integer> iterator) {
        if (iterator.hasNext()) {
            return iterator.next();
        }
        throw new NoSuchElementException();
    }
}